from langchain_core.output_parsers import StrOutputParser
from neo4j import GraphDatabase
import llm_shift

import logging

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

graph = GraphDatabase.driver('bolt://localhost:7687', auth=("neo4j", "12345678"))

model_instance = llm_shift.ChatModel("model_db", "ep-20240612030040-z6dxd", 0)
chosen_model = model_instance.model_choose()
chain = (
        chosen_model
        | StrOutputParser()
)
#result = chain.invoke("你好")

prompt = ("你是一个Cypher查询语句专家，我需要你把中文语言转化成Cypher查询语句在graphdb里建立完整的关系,严格按照样例格式输出，不要有多余的文字。"
          "example：'小王是小李的妈妈，小明是小红的弟弟，小李和小明是同学。'"
          "output：MERGE (:Person {name: '小王'})"
          "MERGE (:Person {name: '小李'})"
          "MERGE (:Person {name: '小明'})"
          "MERGE (:Person {name: '小红'})"
          "MERGE (:Person {name: '小王'})-[:MOTHER]->(:Person {name: '小李'})"
          "MERGE (:Person {name: '小明'})-[:BROTHER]->(:Person {name: '小红'})"
          "MATCH (l1:Person {name: '小李'}), (l2:Person {name: '小明'})"
          "MERGE (l1)-[:CLASSMATE]->(l2)"
          )
prompt1 = ("你是一个Cypher查询语句专家，我需要你根据中文语言转化成Cypher查询语句来检索graphdb中的关系，严格按照样例格式输出，不要有多余的文字。"
           "example：'查找名为 '小李' 的节点与其他节点之间的关系'"
           "output："
           "MATCH (a:Person {name: '小李'})-[r]-(b:Person)"
           "RETURN a, b, type(r)")

text = '''
小张是小李的爸爸，小李是小张和小花的儿子，
小花是小张的妻子，小红是小李的表姐，
因为小红的妈妈是小李的妈妈的妹妹，
小明是小李的同学，同时也是小红的邻居，
小华是小李的叔叔，也就是小张的弟弟,
小刚是小华的儿子，也就是小李的堂弟，
小芳是小刚的女朋友，也是小李的好朋友，
小强是小芳的弟弟，他和小李、小明一起上学，
小美是小李的表亲，因为她的妈妈是小李的爸爸的表妹，
小亮是小美的男朋友，也是小明的远房亲戚，
小丽是小李的邻居，她和小红是好朋友，并且经常一起参加社区活动，
小杰是小丽的哥哥，他和小刚在同一个篮球队打球，
小兰是小杰的姐姐，她和小张是同事，
小青是小兰的女儿，她和小芳是大学同学，
小海是小青的哥哥，他和小明在同一个足球队，
小云是小海的女朋友，她和小李在同一个绘画班学习。
'''


def write():
    try:
        query = chain.invoke(prompt + text)
        logging.info(f"Write query: {query}")
        with graph.session() as session:
            result = session.run(query)
            print(result)
        # 这里可以选择记录日志而不是打印
        logging.info("Write operation completed successfully.")
    except Exception as e:
        logging.error(f"Failed to write: {e}")


def read():
    try:

        query = chain.invoke(prompt1+"我想知道小云与小李的关系。")
        s = ""

        with graph.session() as session:
            result = session.run(query)
            for record in result:
                s = s + str(record)
        logging.info(f"Read query: {query}")
        logging.info(f"Read results: {s}")
        # 使用info日志记录结果
        logging.info(
            chain.invoke(s + "请严格按照example简述结果（example:A是B的同学，A是C的兄弟......）"))
    except Exception as e:
        logging.error(f"Failed to read: {e}")


# 调用函数
write()
read()
